import numpy as np
import matplotlib.pyplot as plt


def entropy(p):
    return p * (-np.log2(p))


def gini(p):
    return p * (1 - p)


p_list = []
entropy_list = []
gini_list = []
for p in np.arange(0, 1, 0.01):
    # 整个事件只有2种可能的取值
    p1 = p
    p2 = 1 - p1
    out = entropy(p1) + entropy(p2)
    out1 = gini(p1) + gini(p2)

    p_list.append(p)
    entropy_list.append(out)
    gini_list.append(out1)

plt.plot(p_list, entropy_list)
plt.plot(p_list, gini_list)
plt.legend(["entropy", "gini"])
plt.show()
